//
// Created by hamlet on 23-3-10.
//

#include "preprocess.h"
#include <filesystem>
#include <fmt/core.h>
#include <fstream>
#include <unordered_map>
#include <unordered_set>
#include "Model.h"
#include "porter.h"

namespace fs = std::filesystem;

// 读取Stop Words列表
std::unordered_set<std::string> readStopWords() {
    std::unordered_set<std::string> stop_words;
    std::ifstream data{"english.stop"};
    if (!data.is_open()) {
        fmt::print("无法打开english.stop文件\n");
    } else {
        std::string word;
        do {
            data >> word;
            stop_words.emplace(word);
        } while (!data.eof());
    }
    return stop_words;
}

const std::unordered_set<std::string> STOP_WORDS{readStopWords()};

// 用函数包装一下以使用比我年龄还大的远古代码
inline void doStem(std::string &word) {
    char str[word.size() + 1];
    strcpy(str, word.c_str());
    str[stem(str, 0, static_cast<int>(word.size()) - 1) + 1] = 0;
    word = str;
}

// 处理单个文件
void processFile(Model *model, const fs::path &path, FansionML::IndexT label_index) {
    std::ifstream data{path};
    if (!data.is_open()) {
        fmt::print("无法打开文件：{}\n", path.c_str());
        return;
    }

    // Step 1
    // 为避免对结果产生影响，忽略文本的头部信息，因为其中包含有文本所属的正确类型
    char c1, c2;
    data.get(c1), data.get(c2);
    while (!(data.eof() || (c1 == '\n' && c2 == '\n'))) {
        c1 = c2;
        data.get(c2);
    }

    auto file_index = model->newFileData(path.filename().c_str(), label_index);

    // Step 2
    // 忽略所有其他符号，只保留英文字符组成的单词，并将大写转换为小写
    // Step 3
    // 搜索停用词列表去除停用词
    // Step 4
    // 使用Porter stemmer算法将词性还原
    // Step 5
    // 将词插入计数列表中
    char c;
    std::string word;
    do {
        data.get(c);
        if ('a' <= c && c <= 'z') {
        } else if ('A' <= c && c <= 'Z') {
            c -= 'A' - 'a';
        } else {
            if (!(word.empty() || STOP_WORDS.contains(word))) {
                doStem(word);
                model->addWord(word, file_index);
            }
            word.clear();
            continue;
        }
        word += c;
    } while (!data.eof());
}

void preprocess(Model *model, const char *str) {
    fs::directory_entry entry{str};
    if (!entry.exists() || !entry.is_directory()) {
        fmt::print("无法打开目录：{}\n", str);
        exit(-1);
    }
    model->clear();
    std::unordered_map<std::string, int> count;
    for (auto &p: fs::directory_iterator(str)) {
        if (p.is_directory()) {
            auto label_index = model->newLabel(p.path().filename().c_str());
            for (auto &pi: fs::directory_iterator(p.path())) {
                if (pi.is_regular_file()) {
                    processFile(model, pi.path(), label_index);
                }
            }
        }
    }
    model->filterWords(5);
    model->transformToEigen();
    model->clear();
    model->calcTfPlusIdf();
}
